import subprocess
import ipaddress
import requests
import json
import urllib3
import os
import shutil

# 关闭警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

def is_private_ip(ip):
    try:
        address = ipaddress.ip_address(ip)
        if isinstance(address, ipaddress.IPv6Address):
            # Exclude link-local and deprecated site-local
            if address.is_link_local or address.is_site_local:
                return True
        return address.is_private
    except ValueError:
        return False

def get_ip_addresses_from_stream(pcap_file, stream_id):
    command = [
        'tshark', 
        '-r', pcap_file, 
        '-Y', f'tcp.stream == {stream_id}', 
        '-T', 'fields', 
        '-e', 'ip.addr',
        '-e', 'ipv6.addr'
    ]

    try:
        output = subprocess.check_output(command, stderr=subprocess.STDOUT).decode('utf-8')  # also redirect stderr to suppress warnings
    except subprocess.CalledProcessError as e:
        print(f"Error during tshark execution: {e}")
        return []

    lines = output.splitlines()
    
    ip_addresses_set = set()
    for line in lines:
        for ip in line.split(','):
            if not is_private_ip(ip.strip()):  # strip to remove leading/trailing whitespace
                ip_addresses_set.add(ip.strip())

    return list(ip_addresses_set)

def get_tcp_stream_ids(pcap_file):
    command = [
        'tshark', 
        '-r', pcap_file, 
        '-T', 'fields', 
        '-e', 'tcp.stream'
    ]

    try:
        output = subprocess.check_output(command).decode('utf-8')
    except subprocess.CalledProcessError as e:
        print(f"Error during tshark execution: {e}")
        return []

    # 解析输出，删除重复的流ID
    stream_ids = list(set(output.split()))

    return stream_ids

def misp_search(tag,page):
    # MISP的URL和API密钥
    misp_url = 'https://192.168.1.187:443/'
    misp_key = 'QirtwTfp4aKuICndvDDNztKGiRGpOpTY62qKtB6r'
    headers = {
        'Authorization': misp_key,
        'Accept': 'application/json',
        'Content-Type': 'application/json'
    }

    # 要搜索的标签
    tag_to_search = tag

    # 构建请求的URL
    search_url = f'{misp_url}/events/restSearch'

    # 构建请求的数据
    data = {
        'returnFormat': 'json',
        'tags': tag_to_search,
        'page' : page,
        'limit': 10
    }

    # 发起POST请求
    response = requests.post(search_url, headers=headers, data=json.dumps(data), verify=False)
    # 检查响应状态码
    if response.status_code == 200:
        # 打印结果
        return response.json()
    else:
        print('Error:', response.status_code)

# 基于属性ID下载文件
def download_attribute_file(attribute_id, path,filename):
    # MISP 详细信息
    MISP_URL = 'https://192.168.1.187:443/'
    ATTRIBUTE_ID = attribute_id
    API_KEY = 'QirtwTfp4aKuICndvDDNztKGiRGpOpTY62qKtB6r'
    # 构建请求 URL
    url = f"{MISP_URL}/attributes/download/{ATTRIBUTE_ID}"
    # 设置请求头
    headers = {
        'Authorization': API_KEY,  
        'Accept':"application/json",
        'Content-Type': "application/json",
    }
    # 发送 GET 请求
    response = requests.get(url, headers=headers , verify=False)
    # 创建完整的文件路径
    full_path = os.path.join(path, filename)
    # 保存文件到本地
    with open(full_path, 'wb') as f:
        f.write(response.content)

def is_malware(event):
    for tag in event['Event']['Tag']:
        if 'secheart:reputation="malware"' in tag['name'] :
            return 1
    return 0
    
def attribute_id(event, value):
    for attribute in event['Event']['Attribute']:
        if value  in attribute['value'] :
            return attribute['value'],attribute['id']
    return 0 , 0

import tarfile

def unpack(log_path):
    with tarfile.open('/root/jx_codes/android_dataset/gz/sandbox-trace-result.tar.gz', 'r:gz') as tar:
            # 提取strace.log文件
            tar.extract('strace.log', path=log_path)

def clean_directory(folder_path):
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"The directory {folder_path} does not exist.")
        return

    # 遍历文件夹中的所有文件和子文件夹
    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        try:
            # 如果是文件夹，则使用shutil.rmtree删除
            if os.path.isdir(file_path):
                shutil.rmtree(file_path)
            # 如果是文件，则使用os.remove删除
            else:
                os.remove(file_path)
        except Exception as e:
            print(f"Failed to delete {file_path}. Reason: {e}")

# 下载文件，并不是生成csv
if __name__ == "__main__":
    tag = 'secheart:analysis="traffic-captured"'
    page =1
    while True :
        a = misp_search(tag,page)
        for event in a['response'] :
            pcap_id = 0
            ui_id = 0
            syscall_id = 0
            malware = 0
            
            malware = is_malware(event)
            print(malware)
            pcap_name , pcap_id=attribute_id(event,'.pcap')
            if pcap_name == 0:
                continue
            md5 = pcap_name.replace('.pcap','')

            if malware == 1 :
                path = '/root/jx_codes/android_dataset/malware/'+ md5
                print(path)
                os.makedirs(path, exist_ok=True)
            else :
                path = '/root/jx_codes/android_dataset/benign/'+ md5
                os.makedirs(path, exist_ok=True)
            ui_name,ui_id = attribute_id(event,'screenshots')
            syscall_name,syscall_id = attribute_id(event,'sandbox-trace-result.tar.gz')

            if ui_id != 0 :
                download_attribute_file(ui_id,path,ui_name)
            
            if pcap_id != 0:
                download_attribute_file(pcap_id,path,pcap_name)

            if syscall_id != 0:
                download_attribute_file(syscall_id,'/root/jx_codes/android_dataset/gz','sandbox-trace-result.tar.gz')
                unpack(path)
                clean_directory('/root/jx_codes/android_dataset/gz')
        page +=1